package probability;

import java.util.Arrays;
import java.util.List;

import basics.Matrix;
import basics.Vector;
import basics.VectorMatrix;

public class MultivariateGaussianDistribution extends ProbabilityDistribution {

	public MultivariateGaussianDistribution(Vector mean, Matrix var) {
		super._params.add(mean);
		super._params.add(var);
		cacheParams();
	}

	public MultivariateGaussianDistribution() {
	}

	public Vector getMu() {
		return (Vector) getParam(0);
	}

	public Matrix getSigma() {
		return (Matrix) getParam(1);
	}

	public void setMu(Vector mu) {
		setParam(0, mu);
	}

	public void setSigma(Matrix sigma) {
		setParam(1, sigma);
	}

	public void setParam(int i, VectorMatrix vm) {
		_params.set(i, vm);
		cacheParams();
	}

	private double _cahed_coeff;
	private Matrix _inv_sigma;

	private void cacheParams() {
		Vector m = getMu();
		m = m.getColumnVector();
		_params.set(0, m);
		Matrix v = getSigma();

		_cahed_coeff = Math.pow(2 * Math.PI, -m.length() / 2) * Math.pow(v.det(), -.5);
		_inv_sigma = v.inverse();
	}

	@Override
	public double measure(Vector x) {
		Vector m = (Vector) getParam(0);
		VectorMatrix x1 = (VectorMatrix) x;
		Vector xminusmean = (Vector) x1.minus(m);
		Vector xminusmeant = (Vector) xminusmean.transpose();

		// it can be further simplified :D
		return _cahed_coeff * Math.exp(xminusmeant.timesEqual(_inv_sigma).timesEqual(xminusmean).times(-.5).get(0, 0));
	}

	@Override
	protected ProbabilityDistribution newInstance() {
		return new MultivariateGaussianDistribution();
	}

	public static MultivariateGaussianDistribution product(MultivariateGaussianDistribution... distributions) {
		return product(Arrays.asList(distributions));
	}

	public static MultivariateGaussianDistribution product(List<MultivariateGaussianDistribution> distributions) {
		MultivariateGaussianDistribution gd = null;
		for (MultivariateGaussianDistribution g : distributions) {
			if (gd == null) {
				gd = new MultivariateGaussianDistribution(g.getMu(), g.getSigma());
			} else {
				Matrix m1 = (Matrix) gd._inv_sigma.add(g._inv_sigma);
				m1 = m1.inverse();
				gd.setSigma(m1);
				Matrix m2 = (Matrix) gd._inv_sigma.times(gd.getMu()).add(g._inv_sigma.times(g.getMu()));
				gd.setMu(m1.times(m2).getColumnVector());
			}
		}

		return gd;
	}

	@Override
	public ProbabilityDistribution times(ProbabilityDistribution distribution) {
		return product(this, (MultivariateGaussianDistribution) distribution);
	}
}
